package com.zzyang.mybatis.jdbc.core;

import com.zzyang.mybatis.jdbc.cluster.Cluster;
import com.zzyang.mybatis.jdbc.cluster.FailFastCluster;
import com.zzyang.mybatis.jdbc.cluster.FailOverCluster;
import com.zzyang.mybatis.jdbc.config.IConstants;
import com.zzyang.mybatis.jdbc.enums.DataSourceType;
import com.zzyang.mybatis.jdbc.spring.SpringContextUtils;
import com.zzyang.mybatis.jdbc.spring.config.SeparationJdbcProperties;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.executor.keygen.SelectKeyGenerator;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.transaction.support.TransactionSynchronizationManager;
import org.springframework.util.StringUtils;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;


/**
 * 拦截器吗,mybatis 每执行一条sql语句都会执行这个拦截器所拦截的方法
 *
 * @author zhuoziyang
 * @date 2021/4/9 7:16
 */
@Intercepts({
        @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class})
})
public class ExecutorInterceptor implements ExtensionInterceptor, InitializingBean {

    private static final String SELECT_KEY_SUFFIX = SelectKeyGenerator.SELECT_KEY_SUFFIX;

    private static final Map<String, DataSourceType> DATA_SOURCE_TYPE_CACHE = new ConcurrentHashMap<>();

    private Cluster cluster;

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        DataSourceType dataSourceType = decidedDataSourceType(invocation);
        //设置数据源
        DataSourceSynchronizationManager.setDataSourceType(dataSourceType);
        try {
            return cluster.execute(invocation);
        } finally {
            DataSourceSynchronizationManager.removeDataSourceType();
            DataSourceSynchronizationManager.removeSlaveDataSourceId();
        }
    }

    @Override
    public Object plugin(Object target) {
        return target instanceof Executor ? Plugin.wrap(target, this) : target;
    }

    private DataSourceType decidedDataSourceType(Invocation invocation) {
        //开启了事务，全走主数据库
        if (TransactionSynchronizationManager.isSynchronizationActive()) {
            return DataSourceType.MASTER;
        }

        //强制走主数据源
        if (DataSourceSynchronizationManager.isForceUseMaster()) {
            return DataSourceType.MASTER;
        }

        Object[] args = invocation.getArgs();
        MappedStatement mappedStatement = (MappedStatement) args[0];

        //该 mappedStatement 被解析过
        DataSourceType dataSourceType = DATA_SOURCE_TYPE_CACHE.get(mappedStatement.getId());
        if (dataSourceType != null) {
            return dataSourceType;
        }

        if (SqlCommandType.SELECT.equals(mappedStatement.getSqlCommandType())) {
            if (mappedStatement.getId().endsWith(SELECT_KEY_SUFFIX)) {
                //这的意思就是查询自增主键的值的意思，要保证跟 insert 在同一个事物中，所以得走主库
                dataSourceType = DataSourceType.MASTER;
            } else {
                dataSourceType = DataSourceType.SLAVE;
            }
        } else {
            dataSourceType = DataSourceType.MASTER;
        }

        DATA_SOURCE_TYPE_CACHE.put(mappedStatement.getId(), dataSourceType);
        return dataSourceType;
    }

    @Override
    public void afterPropertiesSet() throws Exception {
        SeparationJdbcProperties properties = SpringContextUtils.getBean(SeparationJdbcProperties.class);
        String cluster = properties.getCluster();
        if (!StringUtils.hasText(cluster) || IConstants.FAIL_FAST_CLUSTER.equals(cluster)) {
            this.cluster = new FailFastCluster();
        } else if (IConstants.FAIL_OVER_CLUSTER.equals(cluster)) {
            this.cluster = new FailOverCluster();
        } else {
            try {
                this.cluster = SpringContextUtils.getBean(cluster, Cluster.class);
            } catch (Exception e) {
                throw new IllegalArgumentException("无法根据 cluster=" + cluster + " 获取集群容错策略");
            }
        }
    }

}
